from torch import optim
from torch import nn

import torch


# 批次大小
batch_size = 100

# 定义网络结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        '''
            input_size：输入特征大小
            hidden_size：LSTM模型数目
            num_layers：LSTM层数
            batch_first：input 和 output (批次,序列长度，特征向量长度)
            LSTM默认inputs(序列长度,批次，特征向量长度)
        '''
        self.lstm = torch.nn.LSTM(
            input_size=1,
            hidden_size=64,
            num_layers=1,
            batch_first=True
        )
        self.out = torch.nn.Linear(in_features=64, out_features=37)
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, x):
        '''
            LSTM输入格式(batch 批次，
                        seq_len 序列长度，
                        feature 特征向量个数)
        '''
        x = x.view(-1,14,1)

        '''
            虽然batch_first=true 但 h_n 和 c_n的第0个维度是num_layer
            output:[batch,seq_len,hidden_size]
            h_n:[num_layer：LSTM层数,
                batch：批次
                hidden_size：隐藏层层数
                ]
            c_n:[num_layer：LSTM层数,
                batch：批次
                hidden_size：隐藏层层数
                ]
        '''
        output, (h_n, c_n) = self.lstm(x)
        output_in_last_timestep = h_n[-1, :, :]
        x = self.out(output_in_last_timestep)
        x = self.softmax(x)

        return x


LR = 0.001

# 定义模型
model = Net()

# 定义损失函数
mse_loss = nn.CrossEntropyLoss()

# 定义优化器
optimizer = optim.Adam(model.parameters(), LR)

def getdata(filename):
    with open(filename,'r',encoding='utf-8') as f:
        data=[]
        lable=[]
        line=''
        while True:
            data.clear()
            lable.clear()
            for i in range(100):
                line=f.readline()
                if not line:
                    break
                line=line.strip().split(' ')
                line=[int(i) for i in line]
                data.append(line[:-1])
                lable.append(line[-1])
            if not line:
                break
            yield data,lable

def train():
    # 模型的训练状态 Dropout 起作用 一部分神经元不工作
    model.train()
    g=getdata('train3.txt')

    for i in range(1000):
        (inputs,lables)=g.__next__()
        inputs=torch.LongTensor(inputs)
        lables=torch.LongTensor(lables)

        inputs = inputs.long()
        lables=lables.long()

        print(type(inputs))
        print(type(lables))


        # 获得模型的预测结果
        out = model(inputs)

        # 交叉熵代价函数( out(batch,c),lables(batch)  )
        loss = mse_loss(out, lables)

        # 梯度清零
        optimizer.zero_grad()

        # 计算梯度
        loss.backward()

        # 修改权值
        optimizer.step()


def test():
    # 模型的测试状态 Dropout不起作用  全部神经元工作
    model.eval()

    g = getdata('train3.txt')
    correct = 0
    for i in range(100):
        (inputs, lables) = g.__next__()
        inputs = torch.LongTensor(inputs)
        lables = torch.LongTensor(lables)

        # 获得模型的预测结果
        out = model(inputs)

        # 获得最大值，以及最大值所在的位置
        _, predicted = torch.max(out, 1)

        # 预测正确的数量
        correct += (predicted == lables).sum()

    print('Test acc:{0}'.format(correct.item() / 10000))





for epoch in range(10):
    print('epoch', epoch)
    train()
    test()